from langchain_community.utilities import SQLDatabase
from langchain_core.output_parsers import StrOutputParser
from langchain_openai import ChatOpenAI
from langchain.agents import AgentType,create_sql_agent
from langchain_core.prompts import ChatPromptTemplate,MessagesPlaceholder
from langchain_community.agent_toolkits.sql.toolkit import SQLDatabaseToolkit
from dotenv import load_dotenv
import os

# 加载环境变量
load_dotenv()

# 初始化数据库链接
db_uri = (f"mysql+mysqlconnector://"
          f"{os.getenv('MYSQL_USER')}"
          f":{os.getenv('MYSQL_PASSWORD')}"
          f"@{os.getenv('MYSQL_HOST')}"
          f":{os.getenv('MYSQL_PORT')}"
          f"/{os.getenv('MYSQL_DATABASE')}")

db = SQLDatabase.from_uri(db_uri)
# 数据库测试链接
# print(db.run("SELECT * FROM `scores`"))

# llm
llm = ChatOpenAI(
    temperature=0.5,
    model='deepseek-chat',          # 指定使用的模型
    openai_api_key=os.getenv('OPENAI_API_KEY'),         # 替换为有效的API密钥
    openai_api_base='https://api.deepseek.com',  # DeepSeek API端点
    # max_tokens=1024              # 限制生成内容的最大长度
)

toolkit = SQLDatabaseToolkit(db=db, llm=llm)

# 自定义提示词模板（必须包含 agent_scratchpad）
custom_prompt = ChatPromptTemplate.from_messages([
    ("system", """
       你是一个专业的sql助手，请按照以下步骤操作：
        1.根据用户问题生成精确的SQL查询
        2.有统计类的处理时先查询出来后在本地进行运算(不要直接使用SQL语句查询)，
        3.执行查询 如果是修改 修改成功后直接返回"修改成功”
        4.将结果用自然语言解释
        
        示例处理：
            - 平均值：保留两位小数（如88.33分）
            - 总和/计数：整数格式（如175分）
            - 查询结果：`[(Decimal('175.0'),)]` → 张三的总成绩是175分
            - 查询结果：`[(88.333333,)]` → 全班的平均成绩是88.33分
            - 查询结果：`[(90.5,)]` → 全班的平均成绩是90.50分
        
        数据库表结构:{table_info}
    """),
    ("user", "{input}"),
    MessagesPlaceholder(variable_name="agent_scratchpad")
])

# 创建SQL Agent
agent = create_sql_agent(
    llm=llm,
    toolkit=toolkit,
    prompt=custom_prompt,
    output_parser=StrOutputParser,
    agent_type="openai-tools",
    verbose=True,  # 显示详细执行过程
    handle_parsing_errors=True,  # 添加错误处理
    # max_iterations=30,  # 防止无限循环
    # early_stopping_method="generate"  # 添加提前终止
    max_execution_time=60,  # 最大执行时间，单位秒
)

# 执行查询
# 执行查询后修正结果
def chat_request(user_input):
    try:
        response = agent.invoke({"input": user_input})
        print(response["output"])
    except Exception as e:
        print(f"出现错误：{e}")

# 用户输入
user_input = input("请输入问题：")
chat_request(user_input)